// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/task_trace_reader.h>

#include <pollux/common/file/file_systems.h>
#include <pollux/core/plan_node.h>
#include <pollux/exec/trace.h>
#include <pollux/exec/trace_util.h>

namespace kumo::pollux::exec::trace {
    TaskTraceMetadataReader::TaskTraceMetadataReader(
        std::string traceDir,
        memory::MemoryPool *pool)
        : traceDir_(std::move(traceDir)),
          fs_(filesystems::getFileSystem(traceDir_, nullptr)),
          traceFilePath_(getTaskTraceMetaFilePath(traceDir_)),
          pool_(pool),
          metadataObj_(getTaskMetadata(traceFilePath_, fs_)),
          tracePlanNode_(ISerializable::deserialize<core::PlanNode>(
              metadataObj_[TraceTraits::kPlanNodeKey],
              pool_)) {
    }

    std::unordered_map<std::string, std::string>
    TaskTraceMetadataReader::queryConfigs() const {
        std::unordered_map<std::string, std::string> queryConfigs;
        const auto &queryConfigObj = metadataObj_[TraceTraits::kQueryConfigKey];
        for (const auto &[key, value]: queryConfigObj.items()) {
            queryConfigs[key.asString()] = value.asString();
        }
        return queryConfigs;
    }

    std::unordered_map<std::string, std::unordered_map<std::string, std::string> >
    TaskTraceMetadataReader::connectorProperties() const {
        std::unordered_map<std::string, std::unordered_map<std::string, std::string> >
                connectorProperties;
        const auto &connectorPropertiesObj =
                metadataObj_[TraceTraits::kConnectorPropertiesKey];
        for (const auto &[connectorId, configs]: connectorPropertiesObj.items()) {
            const auto connectorIdStr = connectorId.asString();
            connectorProperties[connectorIdStr] = {};
            for (const auto &[key, value]: configs.items()) {
                connectorProperties[connectorIdStr][key.asString()] = value.asString();
            }
        }
        return connectorProperties;
    }

    core::PlanNodePtr TaskTraceMetadataReader::queryPlan() const {
        return tracePlanNode_;
    }

    std::string TaskTraceMetadataReader::nodeName(const std::string &nodeId) const {
        const auto *traceNode = core::PlanNode::findFirstNode(
            tracePlanNode_.get(),
            [&nodeId](const core::PlanNode *node) { return node->id() == nodeId; });
        return std::string(traceNode->name());
    }

    std::string TaskTraceMetadataReader::connectorId(
        const std::string &nodeId) const {
        const auto *traceNode = core::PlanNode::findFirstNode(
            tracePlanNode_.get(),
            [&nodeId](const core::PlanNode *node) { return node->id() == nodeId; });
        const auto *tableScanNode =
                dynamic_cast<const core::TableScanNode *>(traceNode);
        POLLUX_CHECK_NOT_NULL(tableScanNode);
        const auto connectorId = tableScanNode->tableHandle()->connectorId();
        POLLUX_CHECK(!connectorId.empty());
        return connectorId;
    }
} // namespace kumo::pollux::exec::trace
